package com.zyl.websocketdemo.demos.web.config;

import org.springframework.context.annotation.Configuration;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.WebSocketMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import org.springframework.web.socket.server.HandshakeFailureException;
import org.springframework.web.socket.server.HandshakeHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;
import org.springframework.web.socket.server.support.DefaultHandshakeHandler;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author zyl
 * @description
 * @since 2024/6/14
 */
@Configuration
@EnableWebSocket
public class WebsocketConfig implements WebSocketConfigurer {

    public static Map<String, WebSocketSession> sessionMap = new ConcurrentHashMap<>();

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry
                .addHandler(new TextWebSocketHandler() { // WebSocketHandler接口
                    @Override
                    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
                        System.out.println("建立websocket连接后的回调," + session + "," + this + RequestContextHolder.getRequestAttributes());
                        // 存入sessionMap,实际应该关联到用户
                        String name = (String) session.getAttributes().get("name");
                        System.out.println("与: {}建立链接成功" + name);
                        sessionMap.put(name, session);
                    }

                    @Override
                    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
                        System.out.println("收到websocket消息," + session + "," + message+ RequestContextHolder.getRequestAttributes());
                    }

                    @Override
                    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
                        System.out.println("websocket连接关闭后回调" + "," + status + RequestContextHolder.getRequestAttributes());
                    }

                    @Override
                    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
                        System.out.println("连接发生错误被回调");
                    }
                }, "/websocket01")
                .addInterceptors(new HandshakeInterceptor() { // HandshakeInterceptor接口
                    @Override // 这里回传的ServerHttpRequest(实现类为ServletServerHttpRequest)包裹了HttpServletRequest
                    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
                        System.out.println("握手前的回调," + wsHandler + ","+ attributes + "," + RequestContextHolder.getRequestAttributes());
                        // 也可以从session中获取用户信息(也可以使用请求头放token,只要能获取到用户都ok)
                        String name = ((ServletServerHttpRequest) request).getServletRequest().getParameter("name");
                        System.out.println("握手拦截器,获取name: {}"+name);
                        synchronized (this) {
                            if (!sessionMap.containsKey(name)) {
                                attributes.put("name", name);
                                System.out.println("sessionMap中不包含此name:{},允许建立链接"+name);
                                return true;
                            }
                            System.out.println("sessionMap中已包含此name:{}"+name);
                            return false;
                        }
                    }

                    @Override
                    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Exception exception) {
                        System.out.println("握手后回调, " + "," + wsHandler + "," + exception + RequestContextHolder.getRequestAttributes());
                    }
                })
                .setHandshakeHandler(new MyHandShakeHandler()) // HandshakeHandler
                .setAllowedOrigins("*")
        // 开启sockJs,记得把websocketEnabled也给开启了,否则客户端只能使用sockJs,而不能使用websocket
        // .withSockJS()
        // .setWebSocketEnabled(true)
        ;
    }

    static class MyHandShakeHandler implements HandshakeHandler {
        // 借助Spring提供的DefaultHandshakeHandler完成握手(仅仅包裹了一下spring封装的DefaultShakeHandler)
        private DefaultHandshakeHandler defaultHandshakeHandler = new DefaultHandshakeHandler();

        @Override
        public boolean doHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) throws HandshakeFailureException {
            System.out.println("开始握手..." + RequestContextHolder.getRequestAttributes());
            boolean flag = defaultHandshakeHandler.doHandshake(request, response, wsHandler, attributes);
            System.out.println("握手完成..." + flag);
            return flag;
        }
    }

}
